import { correct } from '../utils/correct';
import { Path, PhraseInfo, PrioritySet } from './priority-set';

export type TPinyinDict = Record<string, string>;

export interface IHmmStartDict {
  data: Record<string, number>;
  default: number;
}

export interface IHmmDict {
  data: Record<string, Record<string, number>>;
  default: number;
}

export interface IHmmQueryParams {
  yinJieList: string[];
  maxNum?: number;
  log?: boolean;
  minProb?: number;
}

/**
 * 隐马尔可夫模型（HMM）
 */
export class HiddenMarkovModel {
  private pinyinDict: TPinyinDict;
  private startDict: IHmmStartDict;
  private emissionDict: IHmmDict;
  private transitionDict: IHmmDict;

  constructor(
    pinyinDict: TPinyinDict,
    startDict: IHmmStartDict,
    emissionDict: IHmmDict,
    transitionDict: IHmmDict
  ) {
    this.pinyinDict = pinyinDict;
    this.startDict = startDict;
    this.emissionDict = emissionDict;
    this.transitionDict = transitionDict;
  }

  setPinyinDict = (pinyinDict: Record<string, string>): void => {
    this.pinyinDict = pinyinDict;
  }

  setStartDict = (startDict: IHmmStartDict): void => {
    this.startDict = startDict;
  }

  setEmissionDict = (emissionDict: IHmmDict): void => {
    this.emissionDict = emissionDict;
  }

  setTransitionDict = (transitionDict: IHmmDict): void => {
    this.transitionDict = transitionDict;
  }

  private start = (state: string): number => {
    return this.startDict.data[state] ?? this.startDict.default;
  }

  private emission = (state: string, observation: string): number => {
    return this.emissionDict.data[state]?.[observation] ?? this.emissionDict.default;
  }

  private transition = (fromState: string, toState: string): number => {
    const prodDict = this.transitionDict.data[fromState];

    if (prodDict) {
      return prodDict[toState] ?? prodDict['default'] ?? this.transitionDict.default;
    }

    return this.transitionDict.default;
  }

  private getStates = (observation: string): string[] => {
    return this.pinyinDict[observation]?.split('') ?? [];
  }

  query = ({
    yinJieList,
    maxNum,
    log = false,
    minProb = 3.14e-200
  }: IHmmQueryParams): Path[] => {
    if (!yinJieList.length) {
      return [];
    }

    // 纠错
    const correctedYinJieList = correct(yinJieList);
    const total = correctedYinJieList.length;

    let dp: Map<string, PrioritySet>[] = [new Map()];
    let curObs = correctedYinJieList[0];
    let prevStates = this.getStates(curObs[1]);
    let curStates = prevStates;

    curStates.forEach(curState => {
      const startScore = Math.max(this.start(curState), minProb);
      const emissionScore = Math.max(this.emission(curState, curObs[1]), minProb);
      const score = log ? Math.log(startScore) + Math.log(emissionScore) : startScore * emissionScore;
      const prioritySet = dp[0].get(curState) ?? new PrioritySet(maxNum);
      prioritySet.put(
        new Path(
          [
            new PhraseInfo(
              [curObs[0]],
              [curObs[1]],
              curState
            )
          ],
          score
        )
      );
      dp[0].set(curState, prioritySet);
    });

    for (let i = 1; i < total; i++) {
      curObs = correctedYinJieList[i];
      if (dp.length === 2) {
        dp = [dp.at(-1)!];
      }
      dp.push(new Map());
      prevStates = curStates;
      curStates = this.getStates(curObs[1]);

      curStates.forEach(curState => {
        if (!dp[1].has(curState)) {
          dp[1].set(curState, new PrioritySet(maxNum));
        }

        prevStates.forEach(prevState => {
          const prevSet = dp[0].get(prevState);
          if (prevSet) {
            prevSet.getPaths().forEach(path => {
              const transitionScore = Math.max(this.transition(prevState, curState), minProb);
              const emissionScore = Math.max(this.emission(curState, curObs[1]), minProb);
              const score = log ? path.score + Math.log(transitionScore) + Math.log(emissionScore) : path.score * transitionScore * emissionScore;
              dp[1].get(curState)?.put(
                new Path(
                  [
                    ...path.phraseInfoList,
                    new PhraseInfo(
                      [curObs[0]],
                      [curObs[1]],
                      curState,
                    )
                  ],
                  score
                )
              );
            });
          }
        });
      });
    }

    const result = new PrioritySet(maxNum);

    dp.at(-1)?.forEach(lastSet => {
      lastSet.getPaths().forEach(item => {
        result.put(new Path(item.phraseInfoList, item.score));
      });
    });

    return result.getSortedPaths();
  }
}